# coding: utf-8
""" demo code for back-projection """
# Copyright (c) Benyuan Liu. All Rights Reserved.
# Distributed under the (new) BSD License. See LICENSE.txt for more info.
from __future__ import absolute_import, division, print_function

#import sys
#from PySide6.QtCore import Qt
#from PySide6.QtWidgets import QApplication, QLabel

import zmq
import atexit
import matplotlib.pyplot as plt
import numpy as np
import matplotlib.animation as animation
import pyeit.eit.bp as bp
import pyeit.eit.protocol as protocol
import pyeit.mesh as mesh
from pyeit.eit.fem import EITForward
from pyeit.mesh.wrapper import PyEITAnomaly_Circle

plt.rcParams['toolbar'] = 'none'

context = zmq.Context(2)
socket  = context.socket(zmq.SUB)
socket.connect("tcp://127.0.0.1:8100")
socket.setsockopt(zmq.SUBSCRIBE,b'')

#matplotlib.use('qtagg')

""" 0. build mesh """
n_el = 16  # nb of electrodes
mesh_obj = mesh.create(n_el, h0=0.1)

""" 1. problem setup """

""" 2. FEM forward simulations """
# setup EIT scan conditions
# adjacent stimulation (dist_exc=1), adjacent measures (step_meas=1)
protocol_obj = protocol.create(n_el, dist_exc=1, step_meas=1, parser_meas="std")

# calculate simulated data
fwd = EITForward(mesh_obj, protocol_obj)
v0 = fwd.solve_eit()

""" 3. naive inverse solver using back-projection """
eit = bp.BP(mesh_obj, protocol_obj)
eit.setup(weight="none")
# the normalize for BP when dist_exc>4 should always be True

# extract node, element, alpha
pts = mesh_obj.node
tri = mesh_obj.element

fig, axes = plt.subplots()


ax1=axes
x, y = pts[:, 0], pts[:, 1]
for i, e in enumerate(mesh_obj.el_pos):
    # 获取电极位置的坐标
    ex, ey = x[e], y[e]
    # 计算电极位置相对于中心的角度
    angle = np.arctan2(ey, ex)
    # 设置标注距离中心的半径，略大于电极位置的半径
    r = np.sqrt(ex**2 + ey**2) + 0.05  # 0.05 可根据需要调整
    # 计算标注位置的坐标
    tx = r * np.cos(angle)
    ty = r * np.sin(angle)
    # 添加编号标注
    ax1.text(tx, ty, str(i + 1), color="r", fontsize=9, fontweight="bold",ha='center', va='center')


def update(n):
    global v0
    v1 = np.frombuffer(socket.recv())
    ds = 192.0 * eit.solve(v1, v0, normalize=True)
    im = ax1.tripcolor(pts[:, 0], pts[:, 1], tri, ds)
    ax1.set_aspect('equal')
    return (im,)
anime = animation.FuncAnimation(fig,update,blit=True)
ax1.set_title(r"Reconstituted $\Delta$ Conductivities")
if(True):
    anime.save('ani.gif')
plt.show()
atexit.register(socket.close())